package task

import (
	"aicloud/monitor-apiserver/pkg/db"
	"aicloud/monitor-apiserver/pkg/metric"
	_const "aicloud/monitor-apiserver/pkg/util/const"
	"aicloud/monitor-apiserver/server/response"
	"fmt"
)

func GetOverview() (*response.OverviewResponse, error) {
	idcCount, err := db.DbClient.GetDataCenterCount()
	if err != nil {
		return nil, err
	}
	supplierCount, err := db.DbClient.GetAllSupplierCount("")
	if err != nil {
		return nil, err
	}
	gpuCount, utilization, err := getAllGpuAndUtilization()
	if err != nil {
		return nil, err
	}
	userCount, err := getTotalUser()
	if err != nil {
		return nil, err
	}
	return &response.OverviewResponse{
		IdcCount:       idcCount,
		GpuCount:       int(gpuCount),
		GpuUtilization: utilization,
		UserCount:      int(userCount),
		SupplierCount:  supplierCount,
	}, nil
}

func getTotalUser() (float64, error) {
	userRes, err := metric.Query(_const.TotalUserPromSql)
	if err != nil {
		return 0, err
	}
	return float64(userRes[0].Value), nil

}

func getAllGpuAndUtilization() (float64, float64, error) {
	var (
		allGpuCount, allUtilization float64
	)

	countRes, err := metric.Query(_const.AllGpuCountPromSql)
	if err != nil {
		return 0, 0, err
	}
	allGpuCount = float64(countRes[0].Value)
	utilRes, err := metric.Query(_const.AllGpuUtilizationPromSql)
	if err != nil {
		return 0, 0, err
	}
	allUtilization = float64(utilRes[0].Value)
	return allGpuCount, allUtilization, nil
}

func GetGpuFaults(idcName string) (*response.GpuFaultStatisticsResponse, error) {

	count, faultCount, err := getGpuAndFaults(idcName)
	if err != nil {
		return nil, err
	}
	return &response.GpuFaultStatisticsResponse{
		TotalGpus:       int(count),
		TotalFaultyGpus: int(faultCount),
	}, nil
}

func getGpuAndFaults(idcName string) (float64, float64, error) {
	var (
		gpuCountQuery, gpuFaultsQuery string
	)
	if idcName == "" {
		gpuCountQuery = _const.AllGpuCountPromSql
		gpuFaultsQuery = `gpu_faults_total`
	} else {
		gpuCountQuery = fmt.Sprintf(`sum(kube_node_status_capacity{resource=~"nvidia_com.*", data_center="%s"})`, idcName)
		gpuFaultsQuery = fmt.Sprintf(`gpu_faults_total{data_center="%s"}`, idcName)
	}
	countRes, err := metric.Query(gpuCountQuery)
	if err != nil {
		return 0, 0, err
	}
	faultRes, err := metric.Query(gpuFaultsQuery)
	if err != nil {
		return 0, 0, err
	}
	return float64(countRes[0].Value), float64(faultRes[0].Value), nil
}

func GetTopDistribution(idcName string) (*response.TopModelDistributionResponse, error) {
	var (
		distributionQuery string
		resp              = &response.TopModelDistributionResponse{GPUModels: make([]response.GPUModelCount, 0)}
	)
	if idcName == "" {
		distributionQuery = `topk(3, max by (model_name) (kube_node_status_capacity{resource=~"nvidia_com.*"}))`
		idcName = "all"
	} else {
		distributionQuery = fmt.Sprintf(`topk(3, max by (model_name) (kube_node_status_capacity{resource=~"nvidia_com.*", data_center="%s"}))`, idcName)
	}
	distributionRes, err := metric.Query(distributionQuery)
	if err != nil {
		return nil, err
	}
	for _, sampleMetric := range distributionRes {
		modelName := sampleMetric.Metric["resource"]
		resp.GPUModels = append(resp.GPUModels, response.GPUModelCount{
			ModelName: string(modelName),
			Count:     int(sampleMetric.Value),
		})
	}
	return resp, nil
}

func GetTpo5GpuDistribution(idcName string) (*response.GPUTop5StatisticsResponse, error) {
	var (
		top5Query string
		resp      = &response.GPUTop5StatisticsResponse{GPUModels: make([]response.GPUModelCount, 0)}
	)
	if idcName == "" {
		top5Query = `topk(5, sum(kube_node_status_capacity{resource=~"nvidia_com.*"}) by (resource))`

	} else {
		top5Query = fmt.Sprintf(`topk(5, sum(kube_node_status_capacity{resource=~"nvidia_com.*", data_center="%s"}) by (resource))`, idcName)
	}
	top5Res, err := metric.Query(top5Query)
	if err != nil {
		return nil, err
	}
	for _, topMetric := range top5Res {
		modelName := topMetric.Metric["resource"]
		resp.GPUModels = append(resp.GPUModels, response.GPUModelCount{
			ModelName: string(modelName),
			Count:     int(topMetric.Value),
		})
	}
	return resp, nil
}

func GetOperationTrend(idcName string, req metric.TrendRequest) (*response.OperationTrendResponse, error) {

	var (
		userTrendQuery                 = `user_count_total`
		orderTrendQuery, gpuTrendQuery string
		resp                           = &response.OperationTrendResponse{Series: response.TimeSeriesResponse{Series: make([]response.Series, 0)}}
	)
	if idcName == "" {
		orderTrendQuery = `sum(order_total_count)`
		gpuTrendQuery = `gpu_rental_total`

	} else {
		orderTrendQuery = fmt.Sprintf(`order_total_count{data_center="%s"}`, idcName)
		gpuTrendQuery = fmt.Sprintf(`gpu_rental_total{data_center="%s"}`, idcName)
	}
	userRes, err := metric.QueryRange(userTrendQuery, req)
	if err != nil {
		return nil, err
	}
	gpuRes, err := metric.QueryRange(gpuTrendQuery, req)
	if err != nil {
		return nil, err
	}
	orderRes, err := metric.QueryRange(orderTrendQuery, req)
	if err != nil {
		return nil, err
	}
	series := metric.GetSeries([]metric.SeriesRequest{{Label: "userCount", Samples: userRes[0].Values}, {Label: "orderCount", Samples: orderRes[0].Values}, {Label: "gpuRentalCount", Samples: gpuRes[0].Values}})
	resp.Series.Series = series
	return resp, nil
}

func GetRentalGpuDistribution(idcName string) (*response.GpuDistributionResponse, error) {
	var (
		rentalQuery string
		resp        = &response.GpuDistributionResponse{IDCs: make(map[string]map[string]int)}
	)
	if idcName == "" {
		rentalQuery = `topk(3, gpu_rental_total) by (data_center, model_name)`
	} else {
		rentalQuery = fmt.Sprintf(`topk(3, gpu_rental_total{"data_center=%s"}) by (model_name)`, idcName)
	}
	res, err := metric.Query(rentalQuery)
	if err != nil {
		return nil, err
	}
	for _, idcRes := range res {
		idc := string(idcRes.Metric["data_center"])
		modelName := string(idcRes.Metric["model_name"])
		if (resp.IDCs[idc]) == nil {
			resp.IDCs[idc] = make(map[string]int)
		}
		resp.IDCs[idc][modelName] = int(idcRes.Value)
	}
	return resp, nil
}

func GetGpuRentalRate() (*response.GPURateStatisticsResponse, error) {
	var (
		utilQuery       = `avg(DCGM_FI_DEV_GPU_UTIL) by (data_center)`
		rentedRateQuery = `(sum(kube_node_status_capacity{resource=~"nvidia_com.*"}) - 
         sum(kube_node_status_allocatable{resource=~"nvidia_com.*"})) / sum(kube_node_status_capacity{resource=~"nvidia_com.*"})`
		sumUtil, count float64
	)
	utilRes, err := metric.Query(utilQuery)
	if err != nil {
		return nil, err
	}
	for _, res := range utilRes {
		sumUtil += float64(res.Value)
		count++
	}
	avgUtil := sumUtil / count
	rentedRes, err := metric.Query(rentedRateQuery)
	if err != nil {
		return nil, err
	}
	rentedRate := rentedRes[0].Value
	unRentedRate := 1 - rentedRate
	return &response.GPURateStatisticsResponse{
		UnRentalRate: float64(unRentedRate),
		RentalRate:   float64(rentedRate),
		UsageRate:    avgUtil,
	}, nil
}

func GetAllSuppliers() (*response.OverViewSupplierStatisticsResponse, error) {
	var (
		total, resided, unResided  int
		residedRate, unResidedRate float64
	)
	suppliers, err := db.DbClient.GetAllSuppliers()
	if err != nil {
		return nil, err
	}
	total = len(suppliers)
	for _, supplier := range suppliers {
		if supplier.IsRegistered {
			resided++
		} else {
			unResided++
		}
	}
	residedRate = float64(resided) / float64(total)
	unResidedRate = float64(unResided) / float64(total)
	return &response.OverViewSupplierStatisticsResponse{
		TotalSuppliers:     total,
		ResidedSuppliers:   response.CountRatio{Count: resided, Ratio: residedRate},
		UnResidedSuppliers: response.CountRatio{Count: unResided, Ratio: unResidedRate},
	}, nil

}

func GetIdcOverview(idcName string) (*response.IDCPreviewResponse, error) {

	var (
		nodeQuery      = fmt.Sprintf(`sum(kube_node_info{data_center="%s"})`, idcName)
		gpuQuery       = fmt.Sprintf(`sum(kube_node_status_capacity{resource=~"nvidia_com.*",data_center="%s"}`, idcName)
		rentedGpuQuery = fmt.Sprintf(`sum(kube_node_status_capacity{resource=~"nvidia_com.*",data_center="%s"}) - 
        sum(kube_node_status_allocatable{resource=~"nvidia_com.*",dara_center="%s"})`, idcName, idcName)
		idleGpuQuery   = fmt.Sprintf(`sum(kube_node_status_allocatable{resource=~"nvidia_com.*",data_center="%s"})`, idcName)
		gpuUtilQuery   = fmt.Sprintf(`avg(DCGM_FI_DEV_GPU_UTIL{data_center="%s"})`, idcName)
		gpuFaultsQuery = fmt.Sprintf(`gpu_faults_total{data_center="%s"}/sum(kube_node_status_capacity{resource=~"nvidia_com.*", data_center="%s"})`, idcName, idcName)
	)
	nodeRes, err := metric.Query(nodeQuery)
	if err != nil {
		return nil, err
	}
	gpuRes, err := metric.Query(gpuQuery)
	if err != nil {
		return nil, err
	}
	rentedRes, err := metric.Query(rentedGpuQuery)
	if err != nil {
		return nil, err
	}
	idleRes, err := metric.Query(idleGpuQuery)
	if err != nil {
		return nil, err
	}
	utilQuery, err := metric.Query(gpuUtilQuery)
	if err != nil {
		return nil, err
	}
	faultQuery, err := metric.Query(gpuFaultsQuery)
	if err != nil {
		return nil, err
	}
	count, err := db.DbClient.GetAllSupplierCount(idcName)
	if err != nil {
		return nil, err
	}
	return &response.IDCPreviewResponse{
		NodeCount:          int(nodeRes[0].Value),
		GPUTotal:           int(gpuRes[0].Value),
		GPURented:          int(rentedRes[0].Value),
		GPUIdle:            int(idleRes[0].Value),
		GPUUtilizationRate: float64(utilQuery[0].Value),
		GPUFailureRate:     float64(faultQuery[0].Value),
		SupplierCount:      count,
	}, nil

}
